import platgo as pg
import numpy as np


class MOKP(pg.Problem):

    def __init__(self, D: int = 250, M: int = 2, P: np.ndarray = None, W: np.ndarray = None):
        self.name = "MOKP"
        self.M = M
        self.D = D
        self.type['multi'], self.type['binary'], self.type['large'] = [True] * 3
        self.borders = []
        super().__init__()
        if P is None or W is None:
            self.P = np.random.randint(low=10, high=100, size=(self.M, self.D))
            self.W = np.random.randint(low=10, high=100, size=(self.M, self.D))
        else:
            self.P = P
            self.W = W
            self.M = P.shape[0]
            self.D = P.shape[1]

    def fix_decs(self, pop: pg.Population, method: int = 0) -> None:
        c = np.sum(self.W, axis=1)/2
        rank = np.argsort(np.max(self.P/self.W, axis=0))
        for i in range(len(pop.decs)):
            while np.any(np.dot(self.W, pop.decs[i].T) > c):
                k = np.argwhere(pop.decs[i, rank])[0]
                pop.decs[i, rank[k]] = 0

    def cal_obj(self, pop: pg.Population) -> None:
        pop.objv = np.sum(self.P, axis=1) - np.dot(pop.decs, self.P.T)

    def get_optimal(self) -> np.ndarray:
        # Generate a point for hypervolume calculation
        point = np.sum(self.P, axis=1).T
        return point


if __name__ == "__main__":
    p = np.array([[1,2,3,4,7,9],[6,5,4,3,2,1]])
    print(p.shape)
    w = np.array([[1,2,3,4,5,6],[7,6,5,4,3,1]])
    prob = MOKP(P=p, W=w)
    decs = np.array([[1,0,1,1,1,1],[0,0,1,1,1,1]])
    pop = pg.Population(decs=decs)
    prob.cal_obj(pop)
    print(pop.objv, "objv")
    point = prob.get_optimal()
    print(point, "pf point")
    print(pop.decs, "pop.decs")
    pop.decs = np.where(pop.decs>0, pop.decs, 2)
    print(pop.decs, "pop.decs")
    prob.fix_decs(pop)
    print(pop.decs, "pop.decs")